package com.one.blocks.common.util;

import lombok.experimental.UtilityClass;
import org.springframework.util.CollectionUtils;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Stack;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author <a href="mailto:idler41@163.com">linfuxin</a> created by 2022-02-20 13:31:41
 */
@UtilityClass
public class TreeUtil {

    public static final Integer DEFAULT_ROOT_LEVEL = 0;

    public static final Integer DEFAULT_ROOT_ID = 0;

    public static <T extends TreeNode<T, K>, K extends Serializable> List<T> buildTree(List<T> data, K rootId) {
        return buildTree(data, rootId, null, null);
    }

    public static <T extends TreeNode<T, K>, K extends Serializable> List<T> buildTree(
            List<T> data, K rootId, Function<T, K> getIdFunction, Function<T, K> getParentIdFunction) {
        if (data == null || data.isEmpty()) {
            return Collections.emptyList();
        }
        // 根据parentId分组
        Map<K, List<T>> parentMap = data.stream()
                .collect(Collectors.groupingBy(i -> getParentIdFunction == null ? i.getParentTreeId() : getParentIdFunction.apply(i)));

        // 从根节点开始构建树
        List<T> rootList = parentMap.get(rootId);
        List<T> result = new ArrayList<>(rootList.size());
        for (T t : rootList) {
            result.add(t);
            recursionFillTree(t, parentMap, getIdFunction);
        }
        return result;
    }

    public static <T extends StrictTreeNode<T, K>, K extends Serializable> List<T> buildStrictTree(List<T> data) {
        return buildStrictTree(data, DEFAULT_ROOT_LEVEL, null, null);
    }

    public static <T extends StrictTreeNode<T, K>, K extends Serializable> List<T> buildStrictTree(
            List<T> data, Function<T, K> getIdFunction, Function<T, K> getParentIdFunction) {
        return buildStrictTree(data, DEFAULT_ROOT_LEVEL, getIdFunction, getParentIdFunction);
    }

    public static <T extends StrictTreeNode<T, K>, K> List<T> buildStrictTree(
            List<T> data, Integer rootLevel, Function<T, K> getIdFunction, Function<T, K> getParentIdFunction) {

        if (data == null || data.isEmpty()) {
            return Collections.emptyList();
        }
        Integer firstLevel = rootLevel + 1;
        List<T> result = new ArrayList<>(data.size());
        Integer upLevel = null;
        Map<K, StrictTreeNode<T, K>> upperLevelMap = new HashMap<>(data.size());
        Map<K, StrictTreeNode<T, K>> currentLevelMap = new HashMap<>(data.size());

        for (T treeNode : data) {
            if (upLevel == null) {
                upLevel = rootLevel;
            }

            K id = getIdFunction == null ? treeNode.getTreeId() : getIdFunction.apply(treeNode);
            K parentId = getParentIdFunction == null ? treeNode.getParentTreeId() : getParentIdFunction.apply(treeNode);

            if (firstLevel.equals(treeNode.getLevel())) {
                result.add(treeNode);
                currentLevelMap.put(id, treeNode);
                continue;
            }

            Integer currentLevel = treeNode.getLevel();
            if (upLevel < currentLevel) {
                // 交换upperLevelMap与currentLevelMap引用
                Map<K, StrictTreeNode<T, K>> temp = upperLevelMap;
                upperLevelMap = currentLevelMap;
                currentLevelMap = temp;

                currentLevelMap.clear();
                upLevel = currentLevel;
            }
            currentLevelMap.put(id, treeNode);

            StrictTreeNode<T, K> upperNode = upperLevelMap.get(parentId);
            // treeNode.setParentNode(upperNode);
            if (upperNode == null) {
                throw new RuntimeException(String.format("parentNode is null!currentNode=%s,parentNode=%s", id, parentId));
            }
            List<T> upperChildren = upperNode.getChildren();
            if (upperChildren == null) {
                upperChildren = new ArrayList<>();
                upperNode.setChildren(upperChildren);
            }
            upperChildren.add(treeNode);
        }

        return result;
    }

    public static <T extends TreeNode<T, K>, K> T get(List<T> tree, K k) {
        for (T t : tree) {
            T data = recursionTraversal(t, k);
            if (data != null) {
                return data;
            }
        }
        return null;
    }

    public static <T extends TreeNode<T, K>, K> List<T> findAllChildren(List<T> tree, K k) {
        T node = get(tree, k);
        if (node == null) {
            return Collections.emptyList();
        }
        List<T> data = new ArrayList<>();
        data.add(node);
        recursionChildren(node, data);
        return data;
    }

    public static <T extends TreeNode<T, K>, K extends Serializable> void deptTraversal(T rootNode, DepthTraversalFunction<T, K> depthTraversalFunction) {
        if (rootNode == null) {
            return;
        }
        recursionTraversal(rootNode, new Stack<>(), depthTraversalFunction);
    }

    public static <T extends TreeNode<T, K>, K extends Serializable> void deptTraversal(List<T> treeList, DepthTraversalFunction<T, K> depthTraversalFunction) {
        for (T t : treeList) {
            recursionTraversal(t, new Stack<>(), depthTraversalFunction);
        }
    }

    private static <T extends TreeNode<T, K>, K extends Serializable> void recursionFillTree(T treeNode, Map<K, List<T>> parentMap, Function<T, K> getIdFunction) {
        K treeId = getIdFunction == null ? treeNode.getTreeId() : getIdFunction.apply(treeNode);
        List<T> children = parentMap.get(treeId);
        if (CollectionUtils.isEmpty(children)) {
            return;
        }
        treeNode.setChildren(CollectionUtils.isEmpty(children) ? null : children);
        for (T child : children) {
            recursionFillTree(child, parentMap, getIdFunction);
        }
    }

    private <T extends TreeNode<T, K>, K extends Serializable> void recursionTraversal(T treeNode, Stack<T> stack, DepthTraversalFunction<T, K> depthTraversalFunction) {
        depthTraversalFunction.traversal(treeNode, stack);
        if (CollectionUtils.isEmpty(treeNode.getChildren())) {
            return;
        }
        stack.push(treeNode);
        for (T child : treeNode.getChildren()) {
            recursionTraversal(child, stack, depthTraversalFunction);
        }
        stack.pop();
    }

    private <T extends TreeNode<T, K>, K> T recursionTraversal(T t, K k) {
        if (t.getTreeId().equals(k)) {
            return t;
        }
        if (CollectionUtils.isEmpty(t.getChildren())) {
            return null;
        }
        for (T child : t.getChildren()) {
            T data = recursionTraversal(child, k);
            if (data != null) {
                return data;
            }
        }
        return null;
    }

    private <T extends TreeNode<T, K>, K> void recursionChildren(T node, List<T> data) {
        if (CollectionUtils.isEmpty(node.getChildren())) {
            return;
        }
        for (T child : node.getChildren()) {
            data.add(child);
            recursionChildren(child, data);
        }
    }
}